import streamlit as st
import cv2
from services import face_recog, db
import pandas as pd
import datetime
import sqlite3

st.title("考勤打卡")
st.info("点击按钮开始或停止考勤")

# 初始化数据库
db.init_db()

if "attendance_running" not in st.session_state:
    st.session_state.attendance_running = False
if "signed_in" not in st.session_state:
    st.session_state.signed_in = set()  # 保存今天已签到的员工ID

col1, col2 = st.columns(2)
with col1:
    if st.button("▶️ 开始考勤"):
        st.session_state.attendance_running = True
with col2:
    if st.button("⏹️ 停止考勤"):
        st.session_state.attendance_running = False
        # ✅ 签退所有已签到的人
        for emp_id in st.session_state.signed_in:
            db.save_sign_out(emp_id)
        st.success("✅ 考勤已停止，签退已记录")

frame_slot = st.empty()

if st.session_state.attendance_running:
    recognizer, label_map = face_recog.load_model()
    if recognizer is None:
        st.error("请先训练模型")
        st.session_state.attendance_running = False
    else:
        cap = cv2.VideoCapture(0)
        face_cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
        )

        while st.session_state.attendance_running and cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, 1.3, 5)

            for (x, y, w, h) in faces:
                face_img = gray[y:y + h, x:x + w]
                try:
                    label_id, confidence = recognizer.predict(face_img)
                except:
                    continue

                if confidence < 60:
                    emp_id = label_map.get(label_id, "unknown ID")

                    conn = sqlite3.connect(db.DB_PATH)
                    c = conn.cursor()
                    c.execute("SELECT name, dept FROM 员工 WHERE emp_id=?", (emp_id,))
                    row = c.fetchone()
                    conn.close()

                    if row:
                        name, dept = row
                        if emp_id not in st.session_state.signed_in:
                            ok, msg = db.save_sign_in(emp_id, name, dept)
                            if ok:
                                st.session_state.signed_in.add(emp_id)
                        text = f"{emp_id}"
                    else:
                        text = f"unregistered({emp_id})"
                else:
                    text = "unknown"

                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.putText(
                    frame, text, (x, y - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2
                )

            if len(faces) == 0:
                cv2.putText(
                    frame, "undetected face", (50, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2
                )

            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_slot.image(frame_rgb, channels="RGB")

        cap.release()

# 展示今日考勤
st.subheader("📊 今日考勤记录")
today = datetime.datetime.now().strftime("%Y-%m-%d")
rows = db.get_attendance(today)
if rows:
    df = pd.DataFrame(
        rows, columns=["ID", "工号", "姓名", "部门", "日期", "签到时间", "签退时间", "时长"]
    )
    st.dataframe(df)
    st.download_button(
        "📥 下载今日考勤",
        df.to_csv(index=False).encode("utf-8"),
        file_name=f"attendance_{today}.csv",
        mime="text/csv"
    )
else:
    st.info("今日暂无考勤记录")
